import os
import sys
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(parent_dir)

import json
import ast
from utils.log import logger
from utils.chatglm_utils import ChatGLMUtils

class TimeRecognition:
    def __init__(self, model) -> None:
        self.model = model
        self.prompt = "假设你是一个时间格式转换器，给你一个字面量的文本输入，需要输出格式化后的时间。如果是离散的时间就以字符串表示，字符串之间以逗号隔开，如果是连续的时间范围则以括号将两个边界时间包括进来即可。现在的时间基准是：2023-07。注意不要输出解释性文本，输入：{}；输出：\n"

    def filter(self, literals):
        """
        从literals字段中提取时间维度的相关json
        :param literals: 字面量数据的json列表
        :return: 时间维度相关的literals
        """
        time_literals = []
        for header in literals:
            if header["literal"] == "none" or header['timeLevel'] == "":
                continue
            else:
                time_literals.append(header)
        return time_literals

    def formatter(self, prediction: str):
        """
        对三种时间数据进行格式化：范围时间 (begin_time or none, end_time or none)
                                多点时间 [time1, time2]
                                单点时间 time
        :param prediction: 模型预测的时间
        :return: 格式化后的操作符和条件字符串元组，以列表形式返回: [(操作符1, 条件值1),(操作符2, 条件值2)]
        """
        res = []

        try:
            # 前处理/判断
            # 替换非法字符
            literal = prediction.replace(" ", "")
            literal = literal.replace("，", ",")
            literal = literal.replace("（", "(")
            literal = literal.replace("【", "[")
            literal = literal.replace("）", ")")
            literal = literal.replace("】", "]")

            try:
                times = ast.literal_eval(literal)
                assert type(times[0]) == str and type(times[1]) == str, "time prediction formatting: time data type must be string"
                # 多点时间[]
                if type(times) == list:
                    res.append(('in', str(literal)))
                # 范围时间()
                elif type(times) == tuple:
                    # 假如（）中只有一个值
                    if len(times) == 1:
                        res.append(('==', times[0]))
                    # 正常的范围时间
                    elif len(times) == 2:
                        if times[0] == 'none' and times[1] == 'none':
                            res.append(('==', '2023'))
                        if times[0] != 'none':
                            res.append(('>=', times[0]))
                        if times[1] != 'none':
                            res.append(('<=', times[1]))
            except:
                # 单点时间
                res.append(('==', literal))
            
            return res

        except Exception as err:
            logger.error("time prediction formatting error[{}]: {}".format(type(err), str(err)))
            raise

    def transform(self, query, time_literals):
        """
        格式化时间字面量
        param time_literals: 时间维度相关的literals
        return: conds中关于时间条件的部分
        """
        conds = []
        origin = []
        for header in time_literals:
            literal = header['literal']

            try:
                prediction = self.model.predict(self.prompt.format(str(literal)))
                origin.append({header['name']: prediction})
            except Exception as err:
                logger.error("time prediction error happaned![{}]: {}".format(type(err), str(err)))
                raise
            
            if prediction == 'none':
                continue
            else:
                res = self.formatter(prediction)
                for item in res:
                    operator = item[0]
                    conds.append([header['id'], operator, item[1]])
        
        return {'origin': origin, 'processed': conds}


# model = ChatGLMUtils()
# model.load_model(local_loading=True,
#                  model_path=f"models/time-chatglm-6b-lora.pt")
# if __name__ == "__main__":
#     time_reg = TimeRecognition(model)
#     query = "去年12月合同金额大于五百万的分部有哪些？"
#     literals = [{'member_amount': 0, 'name': '年', 'literal': '去年', 'operator': '==', 'timeLevel': 'year',
#                  'id': 'AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.02de0a83c07e6c06b5f6dc8d09f9fce7-C_DATE_Year-LEVEL-1648634671682',
#                  'type': 'STRING'},
#                 {'member_amount': 0, 'name': '月', 'literal': '去年12月', 'operator': '==', 'timeLevel': 'month',
#                  'id': 'AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.3943e92bed35dac104c3a1f2b2195a5e-LEVEL-1667896190875',
#                  'type': 'STRING'},
#                 {'member_amount': 23, 'name': '销售分部', 'literal': 'none', 'operator': 'none', 'timeLevel': '',
#                  'id': 'AUGMENTED_DATASET_FIELD.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.2b4ff0f606441075f12c1be466e3642e',
#                  'type': 'STRING'},
#                 {'member_amount': 0, 'name': '合同金额', 'literal': '大于五百万', 'operator': '>', 'timeLevel': '',
#                  'id': 'AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8a8ae5ca0178549554951b9501785d983aaa005e',
#                  'type': 'DOUBLE'}]
#     time_literals = time_reg.filter(literals)
#     print(time_reg.transform(query, time_literals))